#!/usr/bin/env python3
import mmap
import os
import struct
import cv2
import numpy as np
import time
from datetime import datetime
from logger import setup_logger
from concurrent.futures import ThreadPoolExecutor

logger = setup_logger()

class SharedMemoryReader:
    """修复后的共享内存读取器，支持异步保存图像"""
    def __init__(self, shm_name="/image_shm", max_size=50 * 1024 * 1024):
        self.shm_name = shm_name
        self.max_size = max_size
        self._open_shared_memory()
        self.set_save_directories()
        # 初始化线程池
        self.executor = ThreadPoolExecutor(max_workers=2)  # 限制为2个线程，适合保存彩色和深度图像
        logger.info("初始化线程池用于异步保存图像")
    
    def set_save_directories(self):
        """创建保存图像的目录"""
        self.color_dir = "./data/Color6"
        self.depth_dir = "./data/Depth6"
        os.makedirs(self.color_dir, exist_ok=True)
        os.makedirs(self.depth_dir, exist_ok=True)
        # logger.info(f"图像保存目录: {self.color_dir}, {self.depth_dir}")
    
    def _open_shared_memory(self):
        """打开共享内存"""
        try:
            self.fd = os.open(self.shm_name, os.O_RDONLY)
            self.shm = mmap.mmap(self.fd, self.max_size, mmap.MAP_SHARED, mmap.PROT_READ)
        except FileNotFoundError:
            logger.error(f"共享内存文件 {self.shm_name} 不存在")
            raise
        
    def read_images(self):
        """读取图像数据，并异步保存"""
        try:
            self.shm.seek(0)
            flag = self.shm.read(5)
            if flag != b'VALID':
                return None, None, None, None

            timestamp = struct.unpack('Q', self.shm.read(8))[0]
            color_size, depth_size = struct.unpack('II', self.shm.read(8))

            # 读取形状 (H, W, C)
            color_shape = struct.unpack('III', self.shm.read(12))
            depth_shape = struct.unpack('III', self.shm.read(12))

            # 读取图像数据
            start_read = time.perf_counter()
            color_img = np.frombuffer(self.shm.read(color_size), dtype=np.uint8).reshape(color_shape)
            depth_img = np.frombuffer(self.shm.read(depth_size), dtype=np.uint16).reshape(depth_shape)
            read_time = time.perf_counter() - start_read
            # logger.info(f"读取图像数据耗时: {read_time:.4f} 秒")

            # 异步提交保存任务
            if color_img is not None and depth_img is not None:
                future = self.executor.submit(self.save_images, str(timestamp), color_img, depth_img)
                # 添加回调函数处理保存结果
                future.add_done_callback(self._save_callback)
            
            return str(timestamp), timestamp, color_img, depth_img
        except Exception as e:
            logger.error(f"读取错误: {e}")
            return None, None, None, None
    
    def _save_callback(self, future):
        """处理异步保存结果的回调函数"""
        try:
            result = future.result()
            success, color_path, depth_path = result
            if success:
                logger.info(f"异步保存图像成功: {color_path}, {depth_path}")
            else:
                logger.error("异步保存图像失败")
        except Exception as e:
            logger.error(f"异步保存图像错误: {e}")
    
    def save_images(self, timestamp, color_img, depth_img):
        """保存图像到文件，并记录每部分耗时"""
        try:
            # 记录总开始时间
            # start_total = time.perf_counter()

            # 1. 保存彩色图像
            # start_color = time.perf_counter()
            color_path = os.path.join(self.color_dir, f"{timestamp}.png")
            cv2.imwrite(color_path, color_img)
            # color_time = time.perf_counter() - start_color
            # logger.info(f"保存彩色图像耗时: {color_time:.4f} 秒, 路径: {color_path}")

            # 2. 深度图像归一化
            # start_normalize = time.perf_counter()
            depth_normalized = cv2.normalize(depth_img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
            # normalize_time = time.perf_counter() - start_normalize
            # logger.info(f"深度图像归一化耗时: {normalize_time:.4f} 秒")

            # 3. 保存深度图像
            # start_depth = time.perf_counter()
            depth_path = os.path.join(self.depth_dir, f"{timestamp}.png")
            cv2.imwrite(depth_path, depth_normalized)
            # depth_time = time.perf_counter() - start_depth
            # logger.info(f"保存深度图像耗时: {depth_time:.4f} 秒, 路径: {depth_path}")

            # 4. 清理旧图片
            # start_cleanup = time.perf_counter()
            self._cleanup_old_images()
            # cleanup_time = time.perf_counter() - start_cleanup
            # logger.info(f"清理旧图片耗时: {cleanup_time:.4f} 秒")

            # 总耗时
            # total_time = time.perf_counter() - start_total
            # logger.info(f"保存图像总耗时: {total_time:.4f} 秒")

            return True, color_path, depth_path
        except Exception as e:
            logger.error(f"保存错误: {e}")
            return False, None, None
    
    def _cleanup_old_images(self):
        """清理旧图片，保持不超过100张"""
        for directory in [self.color_dir, self.depth_dir]:
            try:
                files = [f for f in os.listdir(directory) if f.endswith('.png')]
                if len(files) > 100:
                    files.sort(key=lambda x: os.path.getctime(os.path.join(directory, x)))
                    for file_to_delete in files[:50]:
                        os.remove(os.path.join(directory, file_to_delete))
                    logger.info(f"清理 {directory}: 删除50张旧图片")
            except Exception as e:
                logger.warning(f"清理错误: {e}")
    
    def close(self):
        """关闭共享内存和线程池"""
        try:
            self.shm.close()
            os.close(self.fd)
            self.executor.shutdown(wait=True)  # 等待线程池任务完成
            logger.info("线程池已关闭")
        except:
            pass